# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import logging
import collections
import os
from typing import List, Union

from coin.base.param_util import to_list
import coin.flow.subscriber as flow_subscriber
from coin.strategy.mm.feed import (get_flow_book_builder, get_flow_book_builder_cls)
from coin.strategy.mm.subscription import get_subscription_key
from coin.proto.coin_feed_subscriber_pb2 import FeedSubscriberConfig
from coin.strategy.mm.subscription import (
    FeedSubscriptionRequest,
    get_feed_sub_request,
)
from coin.exchange.base.kr_rest.product import NormProduct


def get_subscription_key_list(product):
  subscription = collections.OrderedDict()
  for p in to_list(product):
    sub_key = get_subscription_key(p)
    if sub_key not in subscription:
      subscription[sub_key] = []
    subscription[sub_key].append(p)
  return subscription


def guess_root_dir(root_dir, endpoint_url):
  if root_dir is None and endpoint_url is None:
    return '/remote/iosg/coin/data/flow'
  elif root_dir is None and endpoint_url is not None:
    raise ValueError(root_dir, endpoint_url)
  elif root_dir is not None and endpoint_url is None:
    assert os.path.exists(root_dir)
    return root_dir
  elif root_dir is not None and endpoint_url is not None:
    assert root_dir.startswith('s3://')
    return root_dir
  else:
    raise ValueError('Never.')


def get_flow_subscriber_from_archive(feed_sub_request,
                                     from_ts,
                                     *,
                                     machine,
                                     root_dir=None,
                                     endpoint_url=None,
                                     worker_id='0',
                                     fs_config=None,
                                     skip_failed_product=False,
                                     sub_req_to_machine_override_map=None):
  root_dir = guess_root_dir(root_dir, endpoint_url)

  topic_string_to_machine_override_map = {}

  topic_string = []
  for sub_key, product_list in feed_sub_request.gen_iter():
    flow_book_builder_cls = get_flow_book_builder_cls(sub_key,
                                                      fs_config=fs_config,
                                                      current_datetime=from_ts)
    for p in product_list:
      try:
        if p.exchange == 'Okex' and p.product_type == 'Futures':  # Special handling for OkexFutures.
          topic_string += flow_book_builder_cls.get_topic_string(p, '1')
          topic_string += flow_book_builder_cls.get_topic_string(p, '2')
        else:
          topic_string += flow_book_builder_cls.get_topic_string(p, worker_id)
        if sub_req_to_machine_override_map and sub_key in sub_req_to_machine_override_map:
          for each_topic_string in topic_string:
            topic_string_to_machine_override_map[
                each_topic_string] = sub_req_to_machine_override_map[sub_key]
      except ValueError as e:
        if skip_failed_product:
          pass
        else:
          raise e

  topic_string = sorted(set(topic_string))
  for ts in topic_string:
    assert '[' not in ts, 'Invalid topic string: %s' % ts

  logging.debug('Topic string: %s', (', '.join(topic_string)))

  flow_sub = flow_subscriber.from_archive(
      from_ts,
      machine=machine,
      root_dir=root_dir,
      endpoint_url=endpoint_url,
      topic_string=topic_string,
      topic_string_to_machine_override_map=topic_string_to_machine_override_map)
  return flow_sub


def run_from_archive(feed_sub_request_or_products: Union[FeedSubscriptionRequest,
                                                         List[NormProduct]],
                     on_book_reset_callbacks,
                     from_ts,
                     to_ts=None,
                     *,
                     machine=None,
                     root_dir=None,
                     endpoint_url=None,
                     worker_id='1',
                     feed_checker_config=None,
                     fs_config=None,
                     use_feed_arbitration=False,
                     skip_failed_product=False,
                     book_builder_opts=None,
                     sub_req_to_machine_override_map=None):
  feed_sub_request = get_feed_sub_request(feed_sub_request_or_products)

  on_book_reset_callbacks = to_list(on_book_reset_callbacks)
  logging.info('Symbols: %s' % (', '.join([str(p) for p in feed_sub_request.all_products])))

  if use_feed_arbitration is not None:
    use_fs_config = FeedSubscriberConfig(use_feed_arbitration=use_feed_arbitration)
    if fs_config is not None:
      fs_config = copy.copy(fs_config)
      fs_config.MergeFrom(use_fs_config)
    else:
      fs_config = use_fs_config

  flow_sub = get_flow_subscriber_from_archive(
      feed_sub_request,
      from_ts,
      machine=machine,
      root_dir=root_dir,
      endpoint_url=endpoint_url,
      worker_id=worker_id,
      fs_config=fs_config,
      skip_failed_product=skip_failed_product,
      sub_req_to_machine_override_map=sub_req_to_machine_override_map)
  for sub_key, product_list in feed_sub_request.gen_iter():
    book_builder = get_flow_book_builder(sub_key=sub_key,
                                         flow_sub=flow_sub,
                                         worker_id=worker_id,
                                         feed_checker_config=feed_checker_config,
                                         fs_config=fs_config,
                                         current_datetime=from_ts,
                                         book_builder_opts=book_builder_opts)
    for callback in on_book_reset_callbacks:
      callback(sub_key, book_builder)

  # Run
  flow_sub.start(until=to_ts)
